import pandas as pd
from sqlalchemy import create_engine, text


def init_engine(sql: str):
    host = 'localhost'
    username = 'root'
    password = '123456'
    database = 'demo'
    port = 3306
    url = f'mysql+pymysql://{username}:{password}@{host}:{port}/{database}'
    print(url)
    engineData = create_engine(
        url,
        pool_recycle=port,
        echo=True,
    )
    connection = engineData.connect()

    # 执行查询并获取结果
    result = connection.execute(text(sql))

    # 获取列名作为表头
    column_names = result.keys()

    # 获取所有数据
    data = result.fetchall()

    # 将数据转换为字典列表，便于创建DataFrame
    data_list = []
    for row in data:
        row_dict = {}
        for i, column_name in enumerate(column_names):
            row_dict[column_name] = row[i]
        data_list.append(row_dict)

    # 创建DataFrame
    df = pd.DataFrame(data_list)

    # 保存到Excel文件
    excel_filename = 'transfer_data.xlsx'
    df.to_excel(excel_filename, index=False)

    print(f"数据已保存到 {excel_filename}")
    print(f"共导出 {len(data_list)} 行数据")
    print(f"表头: {list(column_names)}")

    # 关闭连接
    connection.close()


if __name__ == '__main__':
    execute_sql = "select * from transfer"
    init_engine(execute_sql)
